import numpy as np
import os
import argparse
import sys
import yaml
from scipy.spatial.transform import Rotation as R
import numpy as np
import json

def validate_args(args):
    if not args.y.endswith('.yaml') and not args.y.endswith('.yml'):
        print("错误: -y 参数必须是一个 YAML 文件")
        sys.exit(1)
    if not args.t.endswith('.json'):
        print("错误: -t 参数必须是一个 JSON 文件")
        sys.exit(1)
    if args.o and not args.o.endswith('.yaml') and not args.o.endswith('.yml'):
        print("错误: -o 参数必须是一个 YAML 文件")
        sys.exit(1)

def get_pose_from_yaml(yaml_data, cam_name='midBack_camera'):
    """
        return:
        {
            "position": [x, y, z],
            "left": [x, y, z],
            "forward": [x, y, z]
        }
    """
    camera_list = yaml_data['static_camera_list']
    camera = dict()
    for cam in camera_list:
        if cam['name'] == cam_name:
            camera = cam
            break
    if not camera:
        print(f"错误: 未找到名为 '{cam_name}' 的相机")
        sys.exit(1)
    return camera

def get_yaml_data_from_pose(pose, name, sample_yaml_data):
    """
    pose: dict containing 'position' and 'orientation'
    name: str, name of the camera to update
    sample_yaml_data: original yaml data to copy other fields from
    return: new yaml data with updated camera pose
    """
    new_yaml_data = sample_yaml_data.copy()
    camera_list = new_yaml_data['static_camera_list']
    type_name = ''

    for cam in camera_list:
        if cam['name'] == 'midBack_camera':
            type_name = cam['type']
            break
    for cam in camera_list:
        if cam['name'] == name:
            camera_list.remove(cam)
            break
    new_cam = {
        'name': name,
        'type': type_name,
        'position': pose['position'].tolist(),
        'forward': pose['forward'].tolist(),
        'left': pose['left'].tolist()
    }
    new_yaml_data['static_camera_list'].append(new_cam)
    return new_yaml_data

def normalize_vector(v):
    """归一化向量"""
    norm = np.linalg.norm(v)
    if norm == 0:
        raise ValueError("Cannot normalize zero vector")
    return v / norm

def create_transformation_matrix_from_plf(position, left, forward):
    """
    通过位置和两个向量（左向量和前向量）创建4×4变换矩阵
    :param position: 位置向量 (x, y, z)
    :param left: 左向量
    :param forward: 前向量
    :return: 4×4变换矩阵
    """
    # 计算上向量
    up = np.cross(left, forward)
    
    # 归一化向量
    left = normalize_vector(left)
    forward = normalize_vector(forward)
    up = normalize_vector(up)
    
    # 构建旋转矩阵
    rotation_matrix = np.eye(4)
    rotation_matrix[:3, :3] = np.column_stack((-left, up, forward))
    
    # 构建平移矩阵
    translation_matrix = np.eye(4)
    translation_matrix[:3, 3] = position
    
    # 直接组合旋转矩阵和平移矩阵
    transformation_matrix = rotation_matrix.copy()
    transformation_matrix[:3, 3] = translation_matrix[:3, 3]
    
    return transformation_matrix

def euler_to_rotation_matrix(pitch, yaw, roll):
    """
    将欧拉角（pitch, yaw, roll）转换为旋转矩阵
    :param pitch: 俯仰角（绕x轴旋转）
    :param yaw: 偏航角（绕z轴旋转）
    :param roll: 横滚角（绕y轴旋转）
    :return: 3×3旋转矩阵
    """
    # 将角度转换为弧度
    pitch_rad = np.deg2rad(pitch)
    yaw_rad = np.deg2rad(yaw)
    roll_rad = np.deg2rad(roll)
    
    # 计算绕x轴的旋转矩阵
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(pitch_rad), -np.sin(pitch_rad)],
        [0, np.sin(pitch_rad), np.cos(pitch_rad)]
    ])
    
    # 计算绕y轴的旋转矩阵
    Ry = np.array([
        [np.cos(roll_rad), 0, np.sin(roll_rad)],
        [0, 1, 0],
        [-np.sin(roll_rad), 0, np.cos(roll_rad)]
    ])
    
    # 计算绕z轴的旋转矩阵
    Rz = np.array([
        [np.cos(yaw_rad), -np.sin(yaw_rad), 0],
        [np.sin(yaw_rad), np.cos(yaw_rad), 0],
        [0, 0, 1]
    ])
    
    # 组合旋转矩阵（注意旋转顺序）
    R = np.dot(Rz, np.dot(Ry, Rx))
    
    return R

def create_transformation_matrix_from_translation_and_rotation(translation, rotation):
    """
    通过平移向量和欧拉角创建4×4变换矩阵
    :param translation: 平移向量 (x, y, z)
    :param rotation: 欧拉角 (pitch, yaw, roll)
    :return: 4×4变换矩阵
    """
    # 提取平移向量
    tx, ty, tz = translation['x'], translation['y'], translation['z']
    
    # 提取欧拉角
    pitch, yaw, roll = rotation['pitch'], rotation['yaw'], rotation['roll']
    
    # 计算旋转矩阵
    R = euler_to_rotation_matrix(pitch, yaw, roll)
    
    # 构建4×4变换矩阵
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = [tx, ty, tz]
    
    return T

def extract_pose_info(transformation_matrix):
    """
    从4×4位姿矩阵中提取位置、前向量和左向量
    :param transformation_matrix: 4×4位姿矩阵
    :return: 包含位置、前向量和左向量的字典
    """
    # 提取位置
    position = transformation_matrix[:3, 3]
    
    # 提取前向量（通常是第三列的前三个元素）
    forward = transformation_matrix[:3, 2]
    
    # 提取左向量（通常是第一列的前三个元素）
    left = -transformation_matrix[:3, 0]
    
    # 将结果存储为字典
    pose_info = {
        "position": position,
        "forward": forward,
        "left": left
    }
    
    return pose_info

def compute_real_pose(forward, left, position, transform_data, inv_flag=False):
    """
    计算真实位姿
    :param forward: 前向量
    :param left: 左向量  
    :param position: 位置向量
    :param transform_data: 变换数据字典，包含 transformation 位姿矩阵
    :param inv_flag: 是否计算逆位姿
    :return: 包含位置、前向量和左向量的字典
    """
    # 创建原始位姿的变换矩阵
    trans_mat_1 = create_transformation_matrix_from_plf(position, left, forward)
    
    # 从变换数据中提取平移和旋转
    transformation = transform_data.get('transformation', None)

    # 创建变换矩阵
    trans_mat_2 = transformation if transformation is not None else np.eye(4)
    
    # 组合变换矩阵
    trans_mat = np.dot(trans_mat_1, trans_mat_2)
    
    # 根据inv_flag决定是否计算逆位姿
    if inv_flag:
        trans_mat = np.linalg.inv(trans_mat)
    
    # 提取位姿信息
    real_pose = extract_pose_info(trans_mat)
    
    return real_pose

def transform_yaml(yaml_path, transform_path, output_path=None):
    if output_path is None:
        output_path = yaml_path
    with open(yaml_path, 'r', encoding='utf-8') as f:
        yaml_data = yaml.safe_load(f)
    midBack_pose = get_pose_from_yaml(yaml_data)
    forward = midBack_pose['forward']
    left = midBack_pose['left']
    position = midBack_pose['position']
    if not os.path.exists(transform_path):
        print(f"transform_path not found: {transform_path} perhaps you are in ubuntu and didn't uncomment watch_and_visualize on adjust.py")
        sys.exit(1)
    with open(transform_path, 'r', encoding='utf-8') as f:
        transform_data = json.load(f)

    midBack_real_pose = compute_real_pose(forward, left, position, transform_data, inv_flag=False)
    new_yaml_data = get_yaml_data_from_pose(midBack_real_pose, 'midBack_real_camera', yaml_data)

    with open(output_path, 'w', encoding='utf-8') as f:
        print(f"Saving to {output_path}")
        yaml.safe_dump(new_yaml_data, f, allow_unicode=True, sort_keys=False)

def main():
    usage="""
用法: python cam_yaml.py -y <yaml_file> -t <transform_file> -o <output_file>
说明: 只允许 -y、-t 和 -o 三个参数
    如果有 -o 则输出到指定文件，否则覆盖 -y 文件
"""
    parser = argparse.ArgumentParser(description=usage)
    parser.add_argument('-y', required=True, help='yaml 文件路径')
    parser.add_argument('-t', required=True, help='to_adjust.json 文件路径')
    parser.add_argument('-o', required=False, help='输出路径')
    args, unknown = parser.parse_known_args()
    if unknown:
        print(f"不支持的参数: {' '.join(unknown)}")
        sys.exit(1)
    validate_args(args)

    yaml_path = args.y
    transform_path = args.t
    output_path = args.o if args.o else yaml_path
    transform_yaml(yaml_path, transform_path, output_path)
if __name__ == "__main__":
    main()
